# Import the Time Series library
import statsmodels.tsa.stattools as ts
import pandas as pd
import pprint
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import datetime        as datetime
                                     
# Import Datetime and the Pandas DataReader
from datetime import datetime
from pandas.io.data import DataReader
from numpy import cumsum, log, polyfit, sqrt, std, subtract
from numpy.random import randn
from pandas.stats.api import ols

DATA_DIR = "./data/"

INPUT_VECTOR = ['spy', 'gld', 'vxx', 'wll', 'arex']


def getData(symbol):
  """ Download the Google OHLCV data from 1/1/2000 to 1/1/2013 """
  data = DataReader(symbol, "yahoo", datetime(2000,1,1), datetime(2013,1,1))

def getDataLocaly(file):
  # Open local csv file
  data = pd.read_csv(DATA_DIR + file, index_col=0)
  data.index = pd.to_datetime(data.index)
  return  data

def adfuller(data):
  """ Output the results of the Augmented Dickey-Fuller test for data
       with a lag order value of 1
  """
  return ts.adfuller(data, 1)

def hurst(ts):
	"""Returns the Hurst Exponent of the time series vector ts"""

	# Create the range of lag values
	lags = range(2, 100)

	# Calculate the array of the variances of the lagged differences
	tau = [sqrt(std(subtract(ts[lag:], ts[:-lag]))) for lag in lags]

	# Use a linear fit to estimate the Hurst Exponent
	poly = polyfit(log(lags), log(tau), 1)

	# Return the Hurst exponent from the polyfit output
	return poly[0]*2.0
  
def getResOfPair(in1, in2):
  # calculate hedge ratio
  res = ols(y=in1, x=in2)
  beta_hr = res.beta.x

  # Calculate the residuals of the linear combination
  return in1 - beta_hr*in2
  
def plot_residuals(df):
  months = mdates.MonthLocator()  # every month
  fig, ax = plt.subplots()
  ax.plot(df.index, df["res"], label="Residuals")
  ax.xaxis.set_major_locator(months)
  ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))
  ax.set_xlim(datetime(2012, 1, 1), datetime(2013, 1, 1))
  ax.grid(True)
  fig.autofmt_xdate()

  plt.xlabel('Month/Year')
  plt.ylabel('Price ($)')
  plt.title('Residual Plot')
  plt.legend()

  plt.plot(df["res"])
  plt.show()  
  
if __name__ == "__main__":

#  for item in INPUT_VECTOR:
    # get data
 #   item_data = getDataLocaly(item + '.csv')
  #  item_hurst = hurst(item_data['Adj Close'])
   # print "Hurst(" + str(item) + "):   " + str(item_hurst)
  in1 = getDataLocaly('wll_test.csv')
  in2 = getDataLocaly('arex_test.csv')
  
  df = pd.DataFrame(index=in1.index)
  df['wll'] = in1['Adj Close']
  df['arex'] = in2['Adj Close']
  
  df['res'] = getResOfPair(in1['Adj Close'], in2['Adj Close'])
  
  print(df)
  plot_residuals(df)
  
  item_hurst = hurst(df['res'])

  print "Hurst(wll arex):   " + str(item_hurst)    
  cadf = adfuller(df['res'])
  pprint.pprint(cadf)
  
    